"""elasticsearch 搜索"""
import time
import asyncio
from elasticsearch.helpers import scan


def async_search(client, index, query):
    """异步搜索"""
    res_ = client.async_search.submit(index=index, body=query, track_total_hits=True, timeout='5s')
    async_search_id = res_.get('id')

    while True:
        yield res_['response']
        if res_.get('is_running') is False:  # 已执行完成
            if async_search_id:
                client.async_search.delete(async_search_id)
            break

        time.sleep(5)
        res_ = client.async_search.get(async_search_id)


def scan_search(client, index, query, size=100000):
    """扫描搜索"""
    _count = 0
    batch = scan(client, index=index, query=query, scroll='5m', raise_on_error=False, preserve_order=True, size=5000)
    for item in batch:
        yield item
        _count += 1
        if _count >= size:
            break


if __name__ == '__main__':
    from common.package.elastic_db import es_db
    es_db.init(hosts='http://192.168.101.79:9200')

    def test_async():
        resp = async_search(es_db.client,
                            index='dsm-*',
                            query={'query': {'bool': {'must': []}},
                                   'size': 10,
                                   'sort': {"_time": {"order": "desc"}},
                                   'aggs': {
                                       "timeline": {
                                           "date_histogram": {
                                               "field": '_time',
                                               'calendar_interval': '1d',
                                               "min_doc_count": 0
                                           }
                                       }
                                    }
                                   })
        for item in resp:
            cost = item['took']
            total = item['hits']['total']['value']
            progress = (item['_shards']['successful'] + item['_shards']['skipped'] + item['_shards']['failed']) / item['_shards']['total']
            count = item['hits']['hits']

            print('-' * 20)
            if count:
                print(count)
                print(max([i['_source']['_time'] for i in count]))
                print(min([i['_source']['_time'] for i in count]))

            print('-' * 20)
            print(total, progress, cost)
            print(item.get('aggregations'))


    def test_scan():
        resp = scan_search(es_db.client,
                           index='dsm-*',
                           query={'query': {'bool': {'must': []}}}, size=20)
        for bat in resp:
            pass


    # test_async()
    test_scan()
